package top.yudoge.dogeeditorserverj.web.controller;

import cn.hutool.core.bean.BeanUtil;
import cn.hutool.core.codec.Base64Encoder;
import cn.hutool.core.util.StrUtil;
import cn.hutool.json.JSONUtil;
import dev.langchain4j.data.message.TextContent;
import dev.langchain4j.data.message.UserMessage;
import dev.langchain4j.service.TokenStream;
import lombok.Data;
import lombok.extern.slf4j.Slf4j;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.http.MediaType;
import org.springframework.validation.annotation.Validated;
import org.springframework.web.bind.annotation.PostMapping;
import org.springframework.web.bind.annotation.RequestBody;
import org.springframework.web.bind.annotation.RequestMapping;
import org.springframework.web.bind.annotation.RestController;
import org.springframework.web.servlet.mvc.method.annotation.SseEmitter;
import top.yudoge.dogeeditorserverj.llm.AIAssistant;
import top.yudoge.dogeeditorserverj.llm.LLMStreamConsumer;
import top.yudoge.dogeeditorserverj.llm.RewriteAgent;
import top.yudoge.dogeeditorserverj.transients.query.ChatQuery;
import top.yudoge.dogeeditorserverj.transients.query.RewriteQuery;

import java.io.IOException;
import java.util.ArrayList;
import java.util.List;
import java.util.concurrent.atomic.AtomicReference;

@Slf4j
@RestController
@RequestMapping("/chat")
public class ChatController {

    @Data
    private static class SQLChatQuery {
        private String dataSourceName;

        private List<String> tableSelected;

        private String databaseName;

        private String schemaName;
    }

    @Autowired
    private AIAssistant aiAssistant;

    @Autowired
    private RewriteAgent rewriteAgent;

    @PostMapping(path = "/scratch", produces = MediaType.TEXT_EVENT_STREAM_VALUE)
    public SseEmitter chatFromScratch(@RequestBody RewriteQuery query) {
        SseEmitter emitter = new SseEmitter(1000 * 60 * 60l);
        rewriteAgent.doRewrite(query, new LLMStreamConsumer() {

            StringBuffer buffer =new StringBuffer();
            @Override
            public void onToken(String token) {
                System.out.print(token);
                try {
                    buffer.append(token);
                    if (!token.endsWith("\n") && !token.startsWith("\n")) {
                        emitter.send(SseEmitter.event().name("token").data(token).build());
                        buffer = new StringBuffer();
                    }
                } catch (IOException e) {
                    throw new RuntimeException(e);
                }
            }

            @Override
            public void onFinish() {
                try {
                    emitter.send(SseEmitter.event().name("complete").data("").build());
                    emitter.complete();
                } catch (IOException e) {
                    throw new RuntimeException(e);
                }
            }

            @Override
            public void onError(Throwable throwable) {
                try {
                    emitter.send(SseEmitter.event().name("error").data(throwable).build());
                } catch (IOException e) {
                    throw new RuntimeException(e);
                }
                emitter.completeWithError(throwable);
            }
        });

        return emitter;
    }

    @PostMapping(produces = {MediaType.TEXT_EVENT_STREAM_VALUE})
    public SseEmitter chat(@Validated @RequestBody ChatQuery query) {

        SseEmitter emitter = new SseEmitter(1000 * 60 * 60l);

        SQLChatQuery chatQuery = new SQLChatQuery();
        BeanUtil.copyProperties(query, chatQuery);

        String json = JSONUtil.toJsonStr(chatQuery);

        List<UserMessage> messages = new ArrayList<>();
        UserMessage dbSelectedMessage = UserMessage.from(TextContent.from(
                "User selected is => \n" + json +"\n You should consider the dataSource, dataBase and schema that user selected now!\n\n User request: " + query.getMessage()
        ));
        messages.add(dbSelectedMessage);

        if (!StrUtil.isBlank(query.getUserCode())) {
            UserMessage userCode = UserMessage.from("UserCode: \n" + TextContent.from(query.getMessage()));
            messages.add(userCode);
        }

        UserMessage userTask = UserMessage.from(TextContent.from(query.getMessage()));
        messages.add(userTask);

        AtomicReference<StringBuffer> buffer = new AtomicReference<>(new StringBuffer());

        TokenStream tokenStream = aiAssistant.chat(messages);
        tokenStream.onPartialResponse(partialResponse -> {
            try {
                System.out.print(partialResponse);
                buffer.get().append(partialResponse);
                if (!partialResponse.endsWith("\\")) {
                    emitter.send(SseEmitter.event().name("token").data(Base64Encoder.encode(buffer.get())).build());
                    buffer.set(new StringBuffer());
                }
            } catch (IOException e) {
                throw new RuntimeException(e);
            }
        });

        tokenStream.onError(error -> {
            try {
                emitter.send(SseEmitter.event().name("error").data(error).build());
                emitter.completeWithError(error);
            } catch (IOException e) {
                emitter.completeWithError(e);
            }
        });

        tokenStream.onCompleteResponse(response -> {
            try {
                emitter.send(SseEmitter.event().name("complete").data("").build());
                emitter.complete();
            } catch (IOException e) {
                emitter.completeWithError(e);
            }
        });

        tokenStream.onToolExecuted(toolExecution -> {
            log.debug("Detected tool request => {}", toolExecution.request());
            log.debug("Detectd tool response => {}", toolExecution.result());
        });

        tokenStream.start();

        return emitter;
    }

}
